package com.example.endusation.milvus;


import ai.djl.MalformedModelException;
import ai.djl.inference.Predictor;
import ai.djl.modality.*;
import ai.djl.repository.zoo.Criteria;          // ✅ 正确导入
import ai.djl.repository.zoo.ZooModel;
import ai.djl.repository.zoo.ModelNotFoundException;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.TranslateException;

import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import org.springframework.stereotype.Service;

import java.io.IOException;


@Service
public class EmbeddingService {

    private ZooModel<String, float[]> model;

    @PostConstruct
    public void init() throws IOException, ModelNotFoundException, MalformedModelException {
        System.out.println("开始下载");
        Criteria<String, float[]> criteria = Criteria.builder()
                .setTypes(String.class, float[].class)
                .optModelUrls("djl://ai.djl.huggingface.pytorch/sentence-transformers/all-MiniLM-L6-v2")
                .optEngine("PyTorch")
                .optProgress(new ProgressBar())
                .build();

        this.model = criteria.loadModel();
        System.out.println("✅ Sentence-BERT 模型加载成功！");
    }

    // 生成文本向量
    public float[] getEmbedding(String text) throws TranslateException {
        try (Predictor<String, float[]> predictor = model.newPredictor()) {
            return predictor.predict(text);
        }
    }

    @PreDestroy
    public void destroy() {
        if (model != null) {
            model.close();
        }
    }
}
